package com.rwind.demo;

import com.google.common.collect.ImmutableList;
import org.apache.calcite.adapter.java.JavaTypeFactory;
import org.apache.calcite.adapter.jdbc.JdbcSchema;
import org.apache.calcite.jdbc.*;
import org.apache.calcite.plan.*;
import org.apache.calcite.plan.volcano.VolcanoPlanner;
import org.apache.calcite.prepare.*;
import org.apache.calcite.rel.RelNode;
import org.apache.calcite.rel.logical.LogicalProject;
import org.apache.calcite.rel.logical.LogicalTableScan;
import org.apache.calcite.rex.RexNode;
import org.apache.calcite.schema.SchemaPlus;
import org.apache.calcite.sql.SqlNode;
import org.apache.calcite.sql.SqlWriter;
import org.apache.calcite.sql.parser.SqlParseException;
import org.apache.calcite.sql.parser.SqlParser;
import org.apache.calcite.sql.parser.impl.SqlParserImpl;
import org.apache.calcite.sql.pretty.SqlPrettyWriter;
import org.apache.calcite.sql.validate.SqlConformanceEnum;
import org.apache.calcite.tools.*;
import org.apache.calcite.util.Util;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.SQLException;
import java.util.Properties;
public class CalciteOptimizerExample {

    private static final String POSTGRESQL_SCHEMA = "PUBLIC";
    private static final String MYSQL_SCHEMA = "mysql";


    public void parseAndValidateSQL() throws SQLException, SqlParseException, ValidationException {
        // Build our connection
        Connection connection = DriverManager.getConnection("jdbc:calcite:");

        // Unwrap our connection using the CalciteConnection
        CalciteConnection calciteConnection = connection.unwrap(CalciteConnection.class);

        // Define our parser Configuration
        SqlParser.Config parserConfig = SqlParser.config();

        // Get a pointer to our root schema for our Calcite Connection
        SchemaPlus rootSchema = calciteConnection.getRootSchema();


        // Attach our Postgres Jdbc Datasource to our Root Schema
        //rootSchema.add("EXAMPLESCHEMA", JdbcSchema.create(rootSchema, "EXAMPLESCHEMA", mockDataSource, null, null));

        // Instantiate a data source, this can be autowired in using Spring as well
        DataSource postgresDataSource = JdbcSchema.dataSource(
                "jdbc:postgresql://localhost/db",
                "org.postgresql.Driver", // Change this if you want to use something like MySQL, Oracle, etc.
                "postgres", // username
                "example"   // password
        );

        // Instantiate a data source, this can be autowired in using Spring as well
        DataSource mysqlDataSource = JdbcSchema.dataSource(
                "jdbc:mysql://localhost/db",
                "com.mysql.jdbc.Driver", // Change this if you want to use something like MySQL, Oracle, etc.
                "Username", // username
                "Password"   // password
        );

        // Attach our Postgres Jdbc Datasource to our Root Schema
        rootSchema.add(POSTGRESQL_SCHEMA, JdbcSchema.create(rootSchema, POSTGRESQL_SCHEMA, postgresDataSource, null, null));

        // Attach our MySQL Jdbc Datasource to our Root Schema
        rootSchema.add(MYSQL_SCHEMA, JdbcSchema.create(rootSchema, MYSQL_SCHEMA, mysqlDataSource, null, null));


        // Build a framework config to attach to our Calcite Planners and  Optimizers
        /*FrameworkConfig config = Frameworks.newConfigBuilder()
                .defaultSchema(rootSchema)
                .build();*/
       // RelBuilder rb = RelBuilder.create(config);


        Frameworks.ConfigBuilder config = Frameworks.newConfigBuilder()
                .defaultSchema(rootSchema)
                .parserConfig(parserConfig)
                .context(Contexts.of(calciteConnection.config()));



        Planner planner = Frameworks.getPlanner(config.build());
        SqlNode node = planner.parse("SELECT * FROM EXAMPLESCHEMA.MYTEMPTABLE");

        /*
        HepProgram program = HepProgram.builder().build();
        HepPlanner planner = new HepPlanner(program);

        planner.setRoot(node);

        RelNode optimizedNode = planner.findBestExp();

        final RelShuttle shuttle = new RelHomogeneousShuttle() {
            @Override public RelNode visit(TableScan scan) {
                final RelOptTable table = scan.getTable();
                if (scan instanceof LogicalTableScan && Bindables.BindableTableScan.canHandle(table)) {
                    return Bindables.BindableTableScan.create(scan.getCluster(), table);
                }
                return super.visit(scan);
            }
        };

        optimizedNode = optimizedNode.accept(shuttle);

        final RelRunner runner = connection.unwrap(RelRunner.class);
        PreparedStatement ps = runner.prepare(optimizedNode);

        ps.execute();

        ResultSet resultSet = ps.getResultSet();
        DBTablePrinter.printResultSet(resultSet);
        */

        SqlNode validateNode = planner.validate(node);
        SqlWriter writer = new SqlPrettyWriter();
        validateNode.unparse(writer, 0,0);

        // Print out our formatted SQL to the console
        System.out.println(ImmutableList.of(writer.toSqlString().getSql()));
    }

    public static void main(String[] args) throws SqlParseException, ValidationException, RelConversionException {
        // 1. 创建SchemaPlus对象
        SchemaPlus schema = Frameworks.createRootSchema(true);
        // 2. 将schema中的表和视图加入到SchemaPlus
        // schema.add("emp", new HrSchema().getEmployeeTable());
        // schema.add("dept", new HrSchema().getDepartmentsTable());
        // 3. 创建SqlParser.Config
        SqlParser.Config parserConfig = SqlParser.configBuilder()
                .setParserFactory(SqlParserImpl.FACTORY)
                .setCaseSensitive(true)
                .setConformance(SqlConformanceEnum.DEFAULT)
                .build();
        // 4. 使用Properties对象创建FrameworkConfig
        Properties properties = new Properties();
        properties.setProperty("optiq", "MyOptimizer");
        FrameworkConfig config = Frameworks.newConfigBuilder()
                .defaultSchema(schema)
                .parserConfig(parserConfig)
                .build();
        // 5. 创建RelOptPlanner对象
        JavaTypeFactory typeFactory = new JavaTypeFactoryImpl();
        Planner planner = Frameworks.getPlanner(config);
        // 6. 解析SQL语句
        SqlNode sqlNode = planner.parse("SELECT * FROM emp WHERE deptno = 10");
        // 7. 验证SQL语句
        SqlNode validatedSqlNode = planner.validate(sqlNode);
        // 8. 将SQL语句转换为逻辑查询计划
        RelNode relNode =
                planner.rel(validatedSqlNode).project();
        // 9. 创建优化器对象
        VolcanoPlanner volcanoPlanner = new VolcanoPlanner();
        // 10. 使用优化器进行优化
        RelTraitSet desiredTraits = relNode.getCluster().traitSet()
                .replace(Convention.NONE)
                .replace(getStringLogicalConvention());
        RelNode optimizedRelNode = planner.transform(0, desiredTraits, relNode);
        // 11. 执行优化后的查询计划



    }
    private static RelTrait getStringLogicalConvention() {
        return new Convention.Impl("stringConvention", RelNode.class);
    }
}